//
// aoints.cc
//
// Copyright (C) 2012 Edward Valeiev
//
// Author: Edward Valeev <evaleev@vt.edu>
// Maintainer: EV
//
// This file is part of the SC Toolkit.
//
// The SC Toolkit is free software; you can redistribute it and/or modify
// it under the terms of the GNU Library General Public License as published by
// the Free Software Foundation; either version 2, or (at your option)
// any later version.
//
// The SC Toolkit is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Library General Public License for more details.
//
// You should have received a copy of the GNU Library General Public License
// along with the SC Toolkit; see the file COPYING.LIB.  If not, write to
// the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA.
//
// The U.S. Government is granted a limited license as per AL 91-7.
//
#include <iostream>
#include <iomanip>
#include <util/options/GetLongOpt.h>
#include <mpqcinit.h>
#include <chemistry/molecule/molecule.h>
#include <chemistry/qc/basis/integral.h>
#include <chemistry/qc/basis/split.h>
#include <TiledArray/eigen.h>
#include "Dmm_scf.h"
#include "direct_tile.h"
#include "tiles.h"
#include "Global.h"
#include "integrals/TilePack.h"
#include "integrals/IntegralEnginePool.h"
#include <util/misc/consumableresources.h>
#include "Input.h"
#include "k-means/atom_ordering.h"
#include <ctime>

namespace TA = TiledArray;
using namespace sc;


int try_main(int argc, char** argv) {

    GetLongOpt opt;
    opt.usage("input_filename");

    MPQCInit init(opt,argc,argv);

    int optind = opt.parse(argc, argv);
    std::string inputfile;
    if (argc - optind == 1) {
        inputfile = argv[optind];
    }
    else {
        opt.usage(std::cout);
        return 1;
    }


    sc::Ref<sc::KeyVal> keyval = init.init(inputfile);
    ConsumableResources::get_default_instance()->print_summary();

    std::cout << "copying mad::world" << std::endl;
    madness::World &world = init.madness_world();
    std::cout << "Got mad::world" << std::endl;

    std::cout << "World rank = " << world.rank() << std::endl;

    //double program_time_i = madness::wall_time();



#if 1

    Ref<Molecule> mol_in = get_input_mol(argc, argv, keyval);


    //Get k-means clusters
    AtomOrdering Ordering(mol_in);
    std::size_t clusters = 4; //get_nclusters(inputfile);
    Ordering.atom_clustering(clusters);
    std::vector<std::size_t> index = Ordering.get_center_index(); // get index
    cluster_atoms cluster_ranges = get_cluster_range(index, clusters); //assign atoms in each TA::range
    Ref<Molecule> mol = Ordering.get_optimized_mol(); // Reorder the molecule


    // Getting basis sets
    Ref<GaussianBasisSet> obs = get_obs_basis(mol);
    Ref<GaussianBasisSet> dfbs = get_dfbs_basis(mol);

    // Create Integral Factory
    Ref<Integral> int_factory = get_integral_factory(argc, argv);
    const std::size_t integral_memory = 82000000;
    int_factory->set_storage(integral_memory);

    //Getting soad density guess from mpqc
    //double soad_time0 = madness::wall_time();
    Ref<SymmSCMatrix> P_guess = get_soad_density(mol, obs);
    //world.gop.fence();
    double soad_timef = madness::wall_time();

    // Get tilesizes for the differnt bases.
    std::vector<unsigned int> tilesize = get_tile_size<unsigned int>(obs, index);
    std::vector<unsigned int> dfbs_tilesize =
                    get_tile_size<unsigned int>(dfbs, index);

    // Constructing tranges for the different arrays.
    std::array<TRange1, 2> block_2d = rank2_tiled_range(tilesize);
    std::array<TRange1, 2> block_dfbs = rank2_tiled_range(dfbs_tilesize);
    std::array<TRange1, 3> block_3d = rank3_tiled_range(tilesize,
                                                        dfbs_tilesize);
    std::array<TRange1, 4> block_4d = rank4_tiled_range(tilesize);

    TA::TiledRange trange_2d(block_2d.begin(), block_2d.end());
    TA::TiledRange trange_2body(block_dfbs.begin(), block_dfbs.end());
    TA::TiledRange trange_3d(block_3d.begin(), block_3d.end());
    TA::TiledRange trange_4d(block_4d.begin(), block_4d.end());

    /// Initialize the Arrays that we want. Array takes world and a tiled range.
    ARRAY2 S_AO(world, trange_2d);
    ARRAY2 H_AO(world, trange_2d);
    ARRAY3 Qpq(world, trange_3d); // Holds 3 cent ints contracted with 2 cent ints

    Bases bases;
    bases.push_back(obs);
    bases.push_back(dfbs);

    // compute overlap integrals
    std::cout << "Computing the overlap matrix" << std::endl;
    int_factory->set_basis(obs);
    IntegralEnginePool<OneBodyInt> overlap_pool(int_factory->overlap());
    double overlap_0 = madness::wall_time();
    make_array_integrals<OneBodyInt>(world, S_AO, bases, cluster_ranges,
                                overlap_pool, 1);
    world.gop.fence();
    double overlap_f = madness::wall_time();
    std::cout << "Overlap time = " << overlap_f - overlap_0 << std::endl;

    // Assign S_AO to Eigen
    //Matrix Eigen_S_AO = TA::array_to_eigen(S_AO);

    // compute core Hamiltonian integrals
    std::cout << "Computing H core integrals" << std::endl;
    double hcore_0 = madness::wall_time();
    IntegralEnginePool<OneBodyInt> hcore_pool(int_factory->hcore());
    make_array_integrals<OneBodyInt>(world, H_AO, bases, cluster_ranges,
                                     hcore_pool, 20);
    world.gop.fence();
    double hcore_f = madness::wall_time();
    std::cout << "hcore time = " << hcore_f - hcore_0 << std::endl;

    // Import Density Guess into TA
    int occ = (int(mol->nuclear_charge()))/int(2); /// if NO3- add a 1
    std::cout << "Density Guess Import" << std::endl;
    int nbasis = bases[0]->nbasis();
    Matrix P_temp(nbasis, nbasis);
    for(int i = 0; i < nbasis; ++i){
        for(int j = 0; j < nbasis; ++j){
            P_temp(i,j) = P_guess->get_element(i,j);
        }
    }
    ARRAY2 D_AO(world,trange_2d);
    import_to_TA(world, D_AO, P_temp); //TA::eigen_to_array<ARRAY2>(world, trange_2d, P_temp);
    //import_to_TA(world, D_AO, P_guess);
    world.gop.fence();
    D_AO("i,j") = 0.5 * D_AO("i,j");

// Direct tile stuff
#if 0
    TA::Array<double, 4, direct_tile_4c> eri_ints(world, trange_4d);

    std::cout << "filling direct array" << std::endl;
    IntegralEnginePool<TwoBodyInt> eri_pool(int_factory->electron_repulsion());
    four_center_int_fill(eri_ints, eri_pool, cluster_ranges);
    world.gop.fence();

    
    ARRAY2 G_AO = D_AO("n,m") * ( 2.0 * eri_ints("i,j,m,n") -  eri_ints("i,n,m,j") );
    std::cout << "G_AO" << std::endl;
    world.gop.fence();

    ARRAY2 F_AO = H_AO("i,j") + G_AO("i,j");
    std::cout << "F_AO" << std::endl;
    world.gop.fence();

    std::cout << "entering direct_DMM" << std::endl;
    double energy = Direct_DMM(world, D_AO, S_AO, H_AO, G_AO, F_AO, eri_ints,
                    mol->nuclear_repulsion_energy() );
#endif

// Direct 2 center G test
#if 0
   TA::Array<double, 2, direct_G_2c> G_AO_Test(world, trange_2d);
   direct_G_fill(G_AO_Test, integral_engines, cluster_ranges, D_AO);
   ARRAY2 stupid = G_AO_Test("i,j");
   world.gop.fence();

#endif

//Direct Coulomb test
#if 0
    // Make the intgrals to give to eigen
    std::cout << "Making two center ints" << std::endl;
    Ref<TwoBodyTwoCenterInt> tbtc = int_factory->electron_repulsion2();
    Matrix E_two_body = Matrix::Zero(dfbs->nbasis(), dfbs->nbasis());
    make_eigen_two_center(tbtc, E_two_body);
    tbtc = 0;
    world.gop.fence();

    //Compute the inverse sqrt of the matrix
    std::cout << "Inverting two center ints" << std::endl;
    Matrix TwoBodyI = E_two_body.inverse();
    world.gop.fence();

    // Give sqrt inverse of two body ints to TA
    std::cout << "Converting metric to TA" << std::endl;
    ARRAY2 DF_metric = TA::eigen_to_array<ARRAY2>(world, trange_2body,
                                                  TwoBodyI);
    world.gop.fence();

    double build0 = madness::wall_time();
    TA::Array<double, 3, direct_coulomb_tile_3c> Q(world, trange_3d);
    int_factory->set_basis(obs, obs, dfbs);
    IntegralEnginePool<TwoBodyThreeCenterInt> eri3_pool(int_factory->electron_repulsion3());
    std::cout << "filling J" << std::endl;
    direct_coulomb_3c_fill(Q, eri3_pool, cluster_ranges);

    std::cout << "Making coulomb " << std::endl;
    TA::Array<double, 1> Pvec = DF_metric("P,Q") * ( D_AO("r,s") * Q("r,s,Q") );

    ARRAY2 J = 2 * Q("i,j,P") * Pvec("P");
    world.gop.fence();
    double build1 = madness::wall_time();

    std::cout << "Build time = " << build1 - build0 << std::endl;

#endif

// Regular DF code.
#if 1
    //Begin three and two center ints
    {
        // This will compute 3 center ints because we gave it a rank 3 tensor.
        std::cout << "Making three center ints" << std::endl;
        ARRAY3 Twobody_Three_Center(world, trange_3d);
        int_factory->set_basis(obs, obs, dfbs);
        IntegralEnginePool<TwoBodyThreeCenterInt> eri3_pool(
                        int_factory->electron_repulsion3());
        double three_center_0 = madness::wall_time();
        make_array_integrals<TwoBodyThreeCenterInt>(world, Twobody_Three_Center,
                                                    bases, cluster_ranges,
                                                    eri3_pool, 20);
        world.gop.fence();
        double three_center_f = madness::wall_time();
        std::cout << "Three center ints time = "
                        << three_center_f - three_center_0 << std::endl;

        // Make the intgrals to give to eigen
        std::cout << "Making two center ints" << std::endl;
        int_factory->set_basis(dfbs, dfbs);
        Ref<TwoBodyTwoCenterInt> tbtc = int_factory->electron_repulsion2();
        Matrix E_two_body = Matrix::Zero(dfbs->nbasis(), dfbs->nbasis());
        make_eigen_two_center(tbtc, E_two_body);
        tbtc = 0;
        world.gop.fence();

        //Compute the inverse sqrt of the matrix
        std::cout << "Inverting two center ints" << std::endl;
        Eigen::SelfAdjointEigenSolver<Matrix> estb(E_two_body);
        Matrix isqrtTC = estb.operatorInverseSqrt();
        world.gop.fence();

        // Give sqrt inverse of two body ints to TA
        ARRAY2 Inverse_Sqrt_Two_Body(world, trange_2body);
        std::cout <<"import to TA" << std::endl;
        import_to_TA(world, Inverse_Sqrt_Two_Body, isqrtTC);
        world.gop.fence();
        isqrtTC.resize(0,0);

        //form Qpq

        world.gop.fence();
        std::cout <<"Breaks in TA"<<std::endl;
        Qpq("p,q,X") = Twobody_Three_Center("p,q,y") *
            Inverse_Sqrt_Two_Body("y,X");
        world.gop.fence(); // Fence to get rid of local arrays
        std::cout <<"Escaped TA" << std::endl;
    }

    // Create a fock and two body integrals array
    std::cout << "Nuclear reuplsion energy = " <<
                  mol->nuclear_repulsion_energy() << std::endl;

    ARRAY2 G_AO = 2.0 * ( Qpq("i,j,X") * ( D_AO("n,m") * Qpq("n,m,X") ) ) -
                ( Qpq("i,n,X") * (  D_AO("n,m") * Qpq("m,j,X") ) );


    double energy = TA::expressions::dot( 2.0 * H_AO("i,j") + G_AO("i,j"),
                             D_AO("i,j") );

    std::cout << "Energy guess using Soad for Naph(cc-pVDZ) = " << energy
                 + mol->nuclear_repulsion_energy() << "\n";

    ARRAY2 F_AO = H_AO("i,j") + G_AO("i,j");

    double Final_energy = DF_DMM(world, D_AO, S_AO, H_AO, F_AO, G_AO, Qpq,
                                 mol->nuclear_repulsion_energy());
    world.gop.fence();
#endif

#endif

    return 0;
}

int main(int argc, char** argv){


    try{
        try_main(argc, argv);
    }
    catch(std::exception &e){
        std::cerr << "Exception thrown = " << e.what() << std::endl;
    }
    catch(...){
        std::cerr << "Unknown exception" << std::endl;
    }

    return 0;
}

